import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *
from dataloader.KITTI2015_loader import KITTI2015
from dataloader.SceneFlow_loader import SceneFlow
from dataloader.USVInland_loader import USVInland

# 数据预处理
mean = [0.406, 0.456, 0.485]
std = [0.225, 0.224, 0.229]
normalize = transforms.Normalize(mean, std)# 归一化：减去均值，除以标准差，消除图像间量纲
transform = transforms.Compose([transforms.RandomCrop((height, width)), transforms.ToTensor(), normalize]) # 组合操作：随机裁剪+转换张量+归一化

# 准备数据集
if dataset == 'kitti2015':
    root = sys_root+'/KITTI/KITTI2015/data_scene_flow'
    train_dataset = KITTI2015(root=root, mode='train', transform=transform, occ=True)
    validate_dataset = KITTI2015(root=root, mode='validate', transform=transform, occ=True)
    test_dataset = KITTI2015(root=root, mode='test', transform=transform)

elif dataset == 'sceneflow':
    root = sys_root+'/SceneFlow'
    train_dataset = SceneFlow(root=root, mode='train', transform=transform)
    validate_dataset = SceneFlow(root=root, mode='validate', transform=transform)
    test_dataset = SceneFlow(root=root, mode='test', transform=transform)

elif dataset == 'usvinland':
    if seg: root = sys_root+'/USVInland/Stereo Matching/Segmentation'
    else: root = sys_root+'/USVInland/Stereo Matching/Low_Res_640_320'
    train_dataset = USVInland(root=root, mode='train', transform=transform)
    validate_dataset = USVInland(root=root, mode='validate', transform=transform)
    test_dataset = validate_dataset # 分割数据集8:2

# 数据加载器
train_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, drop_last=True)
validate_loader = data.DataLoader(validate_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, drop_last=True)
test_loader = data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, drop_last=True)
